import os


def read_video_names(datalist_path):
    video_names = []
    with open(datalist_path, "r") as f:
        for line in f:
            content = line.strip()
            if len(content) == 0:
                continue
            video_names.append(content)
    return video_names


def check_data_integrity(dataset_dir, video_names, check_flag):
    for video_name in video_names:
        if check_flag[0] == "1":  # check 2Dseg
            data_dir = os.path.join(dataset_dir, video_name, "2Dseg")
            all_exist = True
            for i in range(300):
                if not os.path.isfile(os.path.join(data_dir, str(i).zfill(5) + ".png")):
                    all_exist = False
                    break
            if not all_exist:
                print(video_name)
            
        if check_flag[2] == "1":  # check obj_init_pose
            data_dir = os.path.join(dataset_dir, video_name, "obj_init_pose")
            all_exist = True
            if not os.path.isfile(os.path.join(data_dir, "00000.pkl")):
                all_exist = False
            if not all_exist:
                print(video_name)
        
        if check_flag[3] == "1":  # check obj_delta_pose
            data_dir = os.path.join(dataset_dir, video_name, "obj_delta_pose")
            all_exist = True
            for i in range(1, 300):
                if not os.path.isfile(os.path.join(data_dir, str(i).zfill(5) + ".pkl")):
                    all_exist = False
                    break
            if not all_exist:
                print(video_name)


if __name__ == "__main__":
    categories = [
        "C2",
        "C5",
        "C7",
        "C12",
    ]
    check_flag = "100000"  # 2Dseg, extrinsics, obj_init_pose, obj_delta_pose, hand_pose, plane
    dataset_dir = "/localdata_sdd/HOI4D_Sim"

    for C in categories:
        video_names = read_video_names(os.path.join("../datalists", C + ".txt"))
        check_data_integrity(dataset_dir, video_names, check_flag)